from fastapi import WebSocket, WebSocketDisconnect
from fastapi import APIRouter
from typing import Dict
from backend.core.heartbeat import update_client_heartbeat
from datetime import datetime

router = APIRouter(tags=["websocket"])

# 存储所有在线WebSocket连接（key：client_id，value：WebSocket对象）
active_connections: Dict[str, WebSocket] = {}


# 广播提醒给所有在线客户端
def broadcast_alert(alert_data: dict):
    for connection in active_connections.values():
        if connection.client_state == "CONNECTED":
            connection.send_json({
                "type": "new_alert",
                "data": alert_data
            })


# 广播提醒关闭指令（一个客户端关闭提醒，所有客户端同步关闭）
def broadcast_alert_close(alert_id: int):
    for connection in active_connections.values():
        if connection.client_state == "CONNECTED":
            connection.send_json({
                "type": "alert_closed",
                "data": {"alert_id": alert_id}
            })


# WebSocket连接端点
@router.websocket("/ws/{client_id}/{client_name}/{client_type}")
async def websocket_endpoint(
        websocket: WebSocket,
        client_id: str,
        client_name: str,
        client_type: str
):
    # 1. 接受客户端连接
    await websocket.accept()
    active_connections[client_id] = websocket
    print(f"客户端 {client_id}（{client_name}）已连接")

    try:
        # 2. 循环接收客户端消息
        while True:
            data = await websocket.receive_json()

            # 处理心跳消息
            if data["type"] == "heartbeat":
                update_client_heartbeat(client_id, client_name, client_type)
                # 回复心跳确认
                await websocket.send_json({
                    "type": "heartbeat_ack",
                    "data": {"timestamp": datetime.now().timestamp()}
                })

            # 下拉所有提醒
            if data["type"] == "get_all_alerts":
                from backend.api.alerts import get_alerts
                from backend.db.database import get_db
                db = next(get_db())
                # 获取提醒（从数据库）
                alerts = get_alerts(is_active=False, db=db)

                # 关键修复：将Alert对象转换为字典列表
                alerts_dict = [alert.to_dict() for alert in alerts]

                # 发送序列化后的提醒数据
                await websocket.send_json({
                    "type": "all_alerts",
                    "timestamp": datetime.utcnow().isoformat(),
                    "data": alerts_dict
                })

            # 处理客户端关闭提醒
            elif data["type"] == "close_alert":
                alert_id = data["data"]["alert_id"]
                broadcast_alert_close(alert_id)  # 广播给所有客户端

            # 处理客户端新增提醒
            elif data["type"] == "create_alert":
                from backend.api.alerts import create_alert
                from backend.db.database import get_db
                db = next(get_db())
                # 调用创建提醒接口
                alert = create_alert(
                    alert=AlertCreate(
                        title=data["data"]["title"],
                        content=data["data"]["content"],
                        alert_time=datetime.fromisoformat(data["data"]["alert_time"]),
                        repeat=data["data"].get("repeat", "none")
                    ),
                    db=db
                )
                # 广播新提醒
                broadcast_alert({
                    "alert_id": alert.id,
                    "title": alert.title,
                    "content": alert.content,
                    "alert_time": alert.alert_time.isoformat(),
                    "repeat": alert.repeat
                })

    except WebSocketDisconnect:
        # 3. 客户端断开连接
        del active_connections[client_id]
        print(f"客户端 {client_id} 已断开连接")